import pandas as pd
from sqlalchemy import create_engine
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
import json
conn = create_engine('mysql+pymysql://root:root123@127.0.0.1:3306/house_data?charset=utf8')
transfer = LabelEncoder()

def getJsonFirst(items):
    try:
        return int(json.loads(items)[-1])
    except:
        return 0

def getData():
    df = pd.read_sql('select * from house_info', con=conn,index_col='id')[:5000]
    X = df[['city', 'rooms_desc', 'area_range', 'houseType','sale_status','price']]
    X['houseType'] = X['houseType'].replace('住宅', 0) \
        .replace('别墅', 1) \
        .replace('商业类', 2) \
        .replace('商业', 3) \
        .replace('酒店式公寓', 4) \
        .replace('底商', 5) \
        .replace('写字楼', 6) \
        .replace('车库', 7)
    X['city'] = transfer.fit_transform(X['city'].values)
    X['rooms_desc'] = X['rooms_desc'].apply(getJsonFirst)
    X['area_range'] = X['area_range'].apply(getJsonFirst)
    X['price'] = X['price'].astype('int')
    return X

def model_train(data):
    # 数据集 测试集划分
    x_train,x_test,y_train,y_test = train_test_split(data[['city','rooms_desc','area_range','houseType','sale_status']],data['price'],test_size=0.25,random_state=1)
    # 创建线性回归对象
    linear2 = LinearRegression()

    # 多元线性回归  模型训练
    linear2.fit(x_train, y_train)

    # 查看截距和系数
    print(linear2.coef_)
    print(linear2.intercept_)
    # 查看拟合效果得分
    print(linear2.score(x_train, y_train))
    # 0.1475478256576207
    return linear2

def pred(model,*args):
    # df[['city', 'rooms_desc', 'area_range', 'houseType', 'sale_status', 'price']]
    city = transfer.transform([args[0]])[0]
    rooms_desc = args[1]
    area_range = args[2]
    houseType = ['住宅','别墅', '商业类', '商业', '酒店式公寓','底商','写字楼','车库'].index(args[3])
    sale_status = args[4]
    print([city,rooms_desc,area_range,houseType,sale_status])
    pred = model.predict([
        [city,rooms_desc,area_range,houseType,sale_status]
    ])
    return pred[0]

if __name__ == "__main__":
    model = model_train(getData())
    print(pred(model,'长沙',5,150,'别墅',3))